# coding=utf-8

from numpy import random
from matplotlib import pyplot as plt

from matplotlib import rcParams
rcParams['text.usetex']     = False
rcParams['font.sans-serif'] = ['Arial']
rcParams['font.serif']      = ['Arial']


random.seed(42)


def votes(group_sizes, vote_probs):
    num_groups = len(group_sizes)
    
    return [random.binomial(group_sizes[i], vote_probs[i])
            if group_sizes[i]
            else 0
            for i in xrange(num_groups)
            ]


def percentage(num, total):
    return float(num)/total*100


def percentages(groups_total, group_shares,  vote_probs):
    group_sizes = groups(groups_total, group_shares)
    
    group_votes = votes(group_sizes, vote_probs)
    votes_total = sum(group_votes)
    
    turnout_percentage = percentage(votes_total, groups_total)  
    vote_percentages   = [percentage(group_votes[i], votes_total)
                          for i in xrange(num_groups)]
    
    return vote_percentages + [turnout_percentage]


def groups(groups_total, group_shares):
    num_groups = len(group_shares)
    
    all_groups_but_last = [int(groups_total*group_shares[i]) 
                           for i in xrange(num_groups)]
    last_group = groups_total - sum(all_groups_but_last)
    
    return all_groups_but_last + [last_group] 
    

def plot(data, num_bins, names):
    vote_data, turnout_data = data[:-1:], data[-1]
    
    num_parties = len(vote_data)
    
    colors_ = ['r', 'g', 'b', 'c', 'm', 'y', 'k']
    colors = colors_[:num_parties:]
    
    fig = plt.figure()
    
    ax1 = fig.add_subplot(2, 2, 2)
    ax2 = fig.add_subplot(2, 2, 4)
    
    ax3 = fig.add_subplot(2, 2, 1)
    ax4 = fig.add_subplot(2, 2, 3)

    ax1.set_xlim(left=0.0, right=100.0)
    ax1.set_ylim(bottom=0.0, auto=True)
    
    ax2.set_ylim(bottom=0.0, auto=True)

    ax3.set_xlim(left=0.0, right=100.0)
    ax3.set_ylim(bottom=0.0, top=100.0)

    ax1.hist(vote_data, 
             num_bins, 
             label=names, 
             histtype='step',
             range=[0.0, 100.0],
             color=colors
             )

    ax2.hist(vote_data[0],
             num_bins, 
             label=names[0], 
             range=[0.0, 100.0],
             color=colors[0]
             )
    
    for vp, name, color in zip(vote_data, names[:-1:], colors):
        ax3.scatter(turnout_data, vp, color=color, label=name, s=4)
        
    ax4.scatter(turnout_data, 
                vote_data[0], 
                color=colors[0], 
                label=names[0], 
                s=4)

    for ax in ax3, ax4:        
        ax.legend(loc='upper right')
        
        ax.set_xlabel(u'% явки')
        ax.set_ylabel(u'% от числа голосовавших')
        
    
    for ax in ax1, ax2:
        ax.legend(loc='upper right')
        
        ax.set_xlabel(u'% от числа голосовавших')
        ax.set_ylabel(u'число участков')
        
    plt.show()


if __name__ == "__main__":
    names = [u'Партия Власти (№1)',
             u'Оппозиция (партия №2)',
             u'Оппозиция (партия №3)',
             u'Оппозиция (партия №4)',
             u'Оппозиция (партия №5)',
             u'Оппозиция (партия №6)',
             u'Oстальные партии (№7)',
             u'Явка'
             ]
    
    groups_total  = 300
     
    num_samples   = 500
    num_bins      = 100*2

    group_shares_and_vote_probs_0 = [[0.49,  0.6],
                                     [0.19,  0.9],
                                     [0.13,  0.9],
                                     [0.12,  0.9],
                                     [0.03,  0.9],
                                     [0.003, 0.9],
                                     [None,  0.9],
                                     ]

    num_groups = len(group_shares_and_vote_probs_0)

    data = [tuple() for _i in xrange(num_groups + 1)] #extra one for turnout

    for group_shares_and_vote_probs in [group_shares_and_vote_probs_0,
                                        #group_shares_and_vote_probs_1
                                        ]:
        group_shares_, vote_probs = zip(*group_shares_and_vote_probs)
        group_shares = group_shares_[:-1:]  
    
        data_piece = zip(*[percentages(groups_total, group_shares, vote_probs) 
                         for i in xrange(num_samples)])
        data = [left + right for left, right in zip(data, data_piece)]
        
    plot(data, num_bins, names)    
